import librosa, datetime
import numpy as np
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False  # 修正负号


# 音频的时间序列
class sound:
    def __init__(self, signal:np.array, sample_rate:int, name: str = ''):
        self.sr = sample_rate
        self.sgn = signal
        self.time = datetime.timedelta(seconds=len(signal) / sample_rate)
        self.name = name

    def __repr__(self):
        return f'<sound>{self.name}[{self.time}]'

    def wave_show(self):
        plt.figure()
        librosa.display.waveshow(self.sgn, sr=self.sr)
        plt.show()

    def get_part(self, start_time: float = 0, end_time: float = None):
        start_idx = round(start_time * self.sr)
        if end_time is not None:
            end_idx = round(end_time * self.sr)
        else:
            end_idx = None
            end_time = self.time.total_seconds()

        return sound(self.sgn[start_idx:end_idx], self.sr, name=f'{self.name}[{round(start_time, 2)}-{round(end_time, 2)}s]')

    def fft(self, start_time: float = 0, end_time: float = None):
        seq = self.get_part(start_time, end_time).sgn

        spectrum = np.fft.fft(seq)  # 通过快速傅里叶变换，计算信号的频谱

        # 添加频率轴
        freq_axis = np.fft.fftfreq(len(seq), d=1 / self.sr)
        freq_axis = freq_axis[:len(freq_axis) // 2]  # 因为对称，所以去一半

        magnitude = np.abs(spectrum[:len(freq_axis)])  # 能量密度

        return tone(list(freq_axis), magnitude, self, self.name)


# 一个音的频谱
class tone:
    def __init__(self, freq: list, magnitude: list, source: sound = None, name: str = ''):
        self.freq = freq
        self.mag = magnitude
        self.source = source
        self.name = name

    def __repr__(self):
        if self.name != '':
            return f'<tone>{self.name}-{self.get_pitch()[0]} << {self.source}'
        else:
            return f'<tone>{self.get_pitch()[0]} << {self.source}'

    def cut_filter(self, threshold: float = None):
        if threshold is None:
            mean = self.get_mean()
            threshold = mean + (self.get_max() - mean) / 4

        freq, amp = [], []
        n = len(self.mag)
        for i in range(n):
            if 16 <= self.freq[i] <= 20000 and self.mag[i] >= threshold:
                freq.append(self.freq[i])
                amp.append(self.mag[i])
        return tone(freq, amp, self.source, self.name)

    def get_main_freq(self, divider=4):
        mean = self.get_mean()
        # 把均值作为基准能量，divider=1时只取最大能量的频率
        return self.cut_filter(threshold=mean + (self.get_max() - mean) / divider)

    def get_main_pitch(self, divider=4):
        return self.get_main_freq(divider).get_pitch()

    def get_freq(self, max_num: int = None) -> list:
        if max_num is None:
            max_num = len(self.mag)

        # 获取排序后的索引值
        sorted_indices = np.argsort(self.mag)
        freq_list, amp_list, i = [], [], 1
        while len(freq_list) < max_num and i < len(self.mag) + 1:
            freq = self.freq[sorted_indices[-i]]
            for f in freq_list:
                semitone_difference = np.log2(f / freq) * 12
                if abs(semitone_difference) < 1:
                    break
            else:
                freq_list.append(freq)
                amp_list.append(self.mag[sorted_indices[-i]])
            i += 1
        return freq_list

    def get_pitch(self, max_num: int = None) -> list:
        freq_list = self.get_freq(max_num)
        res = []
        for f in freq_list:
            main_pitch = tone.freq2pitch(f)
            res.append(main_pitch)
        return res

    @staticmethod
    def freq2pitch(frequency: list | float):
        def f2p(freq):
            if freq > 16:
                # 音符名称映射表
                note_names = {
                    0: 'C',
                    1: 'C#',
                    2: 'D',
                    3: 'D#',
                    4: 'E',
                    5: 'F',
                    6: 'F#',
                    7: 'G',
                    8: 'G#',
                    9: 'A',
                    10: 'A#',
                    11: 'B'
                }

                # 公式：log2(freq/440) = (x-57)/12, 令C0为第一个音，x表示该音为第几个音
                idx = round(np.log2(freq/440)*12 + 57)
                return f"{note_names[idx % 12]}{idx // 12}"
            else:
                raise ValueError(f'frequency {freq} not > C0 (16hz)')

        if isinstance(frequency, list):
            res = []
            for freq in frequency:
                res.append(f2p(freq))
            return res
        else:
            return f2p(frequency)

    def get_peek(self):
        peek_freq_list, peek_amp_list, increasing = [], [], False
        for i in range(len(self.mag) - 1):
            if self.mag[i] <= self.mag[i + 1]:
                increasing = True
            else:
                if increasing:
                    increasing = False
                    peek_freq_list.append(self.freq[i])
                    peek_amp_list.append(self.mag[i])

        return tone(peek_freq_list, peek_amp_list)

    def get_mean(self):
        return np.mean(self.mag)

    def get_max(self):
        return np.max(self.mag)

    def show(self):
        if self.source:
            fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8))
            fig.suptitle(f"{self.name}")
            ax1.plot(self.freq, self.mag)
            ax1.set_xlabel("频率(hz)")
            ax1.set_ylabel("能量密度")

            # 生成时间轴
            total_time = len(self.source.sgn) / self.source.sr  # 计算总时间
            time_axis = np.linspace(0, total_time, num=len(self.source.sgn))  # 生成时间轴

            ax2.plot(time_axis, self.source.sgn)
            ax2.set_xlabel('时间(s)')
            ax2.set_ylabel('幅度(dB)')
            plt.show()
        else:
            plt.figure(figsize=(10, 8))
            plt.title(f"{self.name}")
            plt.plot(self.freq, self.mag)
            plt.xlabel("频率(hz)")
            plt.ylabel("幅度(dB)")
            plt.show()

    @staticmethod
    def generate(time=2):
        pass


if __name__ == "__main__":
    # 测试音高转换
    # print(tone.freq2pitch([16.5, 261.6, 293.4, 329.6, 349.2, 392, 440, 493.9, 523.3]))

    audio_path = '../music sample/花海伴奏节选.mp3'
    # audio_path = '../music sample/dzy音阶.flac'
    signal, sample_rate = librosa.load(audio_path, sr=None, mono=False)

    # 左右声道
    s_l = sound(signal[0], sample_rate, name="origin")
    s_r = sound(signal[1], sample_rate)
    s_l.wave_show()

    # print(s_l.fft().get_main_pitch())

    time_delta = 0.05
    t1, t2 = 0.2, 2
    s_l.get_part(t1, t2).wave_show()
    sep = list(np.arange(t1, t2 + time_delta, time_delta))

    for i in sep:
        s = s_l.get_part(start_time=i, end_time=i + time_delta)
        t = s.fft()
        f = t.cut_filter(8)
        # f.show()
        print(f.get_pitch(3))

    # A5--B5--C#6--E6-E6-, C#6--B5--A5--B5-C#6-
